import pandas as pd
from sklearn.feature_selection import VarianceThreshold
from common_import import *


def categorize_features(X, threshold=5):
    """
    将特征分为连续型和离散型。

    参数:
    X -- DataFrame, 输入的特征数据
    threshold -- int, 用于将取值少于或等于阈值的数值型特征标识为类别型

    返回:
    continuous_columns -- list, 连续型特征的列名列表
    categorical_columns -- list, 类别型特征的列名列表
    """
    continuous_columns = []
    categorical_columns = []

    for column in X.columns:
        unique_values = X[column].nunique()
        if unique_values <= threshold:
            categorical_columns.append(column)
        else:
            continuous_columns.append(column)

    return continuous_columns, categorical_columns


def variance_filter(X, continuous_columns, threshold=0.25, plot=False):
    """
    使用方差过滤器筛选连续型数据，留下方差大于指定阈值的特征，并绘制筛选掉数据数量随阈值变化的直方图。
    """
    selector = VarianceThreshold(threshold=threshold)
    X_continuous = X[continuous_columns]
    X_filtered_continuous = selector.fit_transform(X_continuous)
    X_filtered_continuous = pd.DataFrame(
        X_filtered_continuous, columns=X_continuous.columns[selector.get_support()]
    )

    # 如果需要绘制直方图
    if plot:
        thresholds = np.arange(0.05, 1.05, 0.05)  # 设定一组阈值从0到1，步长为0.05
        dropped_counts = []

        for t in thresholds:
            selector = VarianceThreshold(threshold=t)
            selector.fit(X_continuous)
            # dropped_count = len(continuous_columns) - sum(selector.get_support())
            dropped_count = sum(selector.get_support())
            dropped_counts.append(dropped_count)

        plt.figure(figsize=(10, 6))
        default_color = plt.rcParams["axes.prop_cycle"].by_key()["color"][0]
        colors = ["red" if t == 0.25 else default_color for t in thresholds]
        plt.bar(thresholds, dropped_counts, width=0.04, alpha=0.7, color=colors)
        plt.xlabel("方差阈值")
        plt.ylabel("筛选留下的特征数量")
        plt.title("")
        plt.xticks(thresholds)
        # plt.grid(True)
        tool.show_or_print("方差筛选.png")
    return X_filtered_continuous


def unique_ratio_filter(X, categorical_columns, threshold=0.05, plot=False):
    """
    筛选类别型数据，留下异众比率（非众数的数量占比）大于指定阈值的特征，并绘制筛选掉数据数量随阈值变化的直方图。
    """
    X_categorical = X[categorical_columns]

    # 初始化列表保存每个阈值对应的筛选掉的特征数量

    if plot:
        thresholds = np.arange(0, 1.05, 0.05)
        dropped_counts = []
        for t in thresholds:
            keep_columns = []
            for column in X_categorical.columns:
                most_frequent_ratio = (
                    X_categorical[column].value_counts(normalize=True).max()
                )
                unique_ratio = 1 - most_frequent_ratio
                if unique_ratio >= t:
                    keep_columns.append(column)
            dropped_count = len(keep_columns)
            dropped_counts.append(dropped_count)
        plt.figure(figsize=(10, 6))
        default_color = plt.rcParams["axes.prop_cycle"].by_key()["color"][0]
        colors = ["red" if t == 0.05 else default_color for t in thresholds]
        plt.bar(thresholds, dropped_counts, width=0.04, alpha=0.7, color=colors)
        plt.xlabel("异众比率阈值")
        plt.ylabel("筛选留下的特征数量")

        plt.xticks(thresholds)
        # plt.show()
        tool.show_or_print("异众比率筛选.png")

    # 筛选并返回最终的过滤后的数据
    keep_columns = []
    for column in X_categorical.columns:
        most_frequent_ratio = X_categorical[column].value_counts(normalize=True).max()
        unique_ratio = 1 - most_frequent_ratio
        if unique_ratio >= threshold:
            keep_columns.append(column)

    X_filtered_categorical = X_categorical[keep_columns]

    return X_filtered_categorical


from scipy.stats import spearmanr


def filter_continuous_with_continuous(X, y, threshold=0.15, plot=False):
    """
    根据斯皮尔曼相关系数筛选连续型特征与连续型目标变量，并绘制筛选留下特征的直方图。
    """
    thresholds = np.arange(0, 1.05, 0.05)  # 生成从0到1，步长为0.05的阈值序列
    retained_counts = []

    if plot:
        for t in thresholds:
            selected_columns = []
            for column in X.columns:
                corr, _ = spearmanr(X[column], y)
                if abs(corr) >= t:
                    selected_columns.append(column)
            retained_counts.append(len(selected_columns))
        plt.figure(figsize=(10, 6))
        default_color = plt.rcParams["axes.prop_cycle"].by_key()["color"][0]
        colors = ["red" if t > 0.14 and t < 0.16 else default_color for t in thresholds]
        plt.bar(thresholds, retained_counts, width=0.04, alpha=0.7, color=colors)
        plt.xlabel("斯皮尔曼相关系数阈值")
        plt.ylabel("筛选留下的特征数量")
        plt.title("")
        plt.xticks(thresholds)
        # plt.show()
        tool.show_or_print("相关性筛选.png")
    selected_columns = []
    for column in X.columns:
        corr, _ = spearmanr(X[column], y)
        if abs(corr) >= threshold:
            selected_columns.append(column)

    return selected_columns


def advanced_filter_features(X, y):
    """
    先通过方差和异众比率筛选特征，然后根据特征类型和目标变量类型进一步筛选。

    参数:
    X -- DataFrame, 输入的特征数据
    y -- Series, 连续型目标变量
    categorical_threshold -- int, 用于标识类别型特征的唯一值数量阈值

    返回:
    X_final_filtered -- DataFrame, 最终筛选后的特征数据
    """
    # 区分连续型和类别型特征
    continuous_columns, categorical_columns = categorize_features(X, threshold=10)
    # print(len(continuous_columns), len(categorical_columns))
    print(
        f"Identified {len(continuous_columns)} continuous features and {len(categorical_columns)} categorical features."
    )

    # 筛选连续型特征
    X_filtered_continuous = variance_filter(X, continuous_columns)
    X_filtered_categorical = unique_ratio_filter(X, categorical_columns)
    print(len(X_filtered_continuous.columns), len(X_filtered_categorical.columns))
    # for column in X_filtered_categorical.columns:
    #     max_ratio = X[column].value_counts(normalize=True).max()
    #     print(f"{column}: {max_ratio},{X[column].nunique()}")

    # print(len(X_filtered_continuous.columns))

    # 根据类别型特征的唯一值数量区分二分类和多分类特征

    # 筛选连续型特征与连续型目标变量
    X_filtered = pd.concat([X_filtered_continuous, X_filtered_categorical], axis=1)
    final_selected_columns = filter_continuous_with_continuous(X_filtered, y)
    # selected_continuous_columns = filter_continuous_with_continuous(
    #     X_filtered_continuous, y
    # )
    # selected_categorical_columns = filter_continuous_with_continuous(
    #     X_filtered_categorical, y
    # )
    # print(len(selected_continuous_columns), len(selected_categorical_columns))
    # final_selected_columns = selected_continuous_columns + selected_categorical_columns

    # for column in selected_categorical_columns:
    #     max_ratio = X[column].value_counts(normalize=True).max()
    #     print(f"{column}: {max_ratio},{X[column].nunique()}")

    # 返回最终筛选后的特征数据
    X_final_filtered = X[final_selected_columns]

    print(f"Final selected features: {len(final_selected_columns)}")

    return X_final_filtered


if __name__ == "__main__":

    # 使用该函数进行特征筛选
    molecular_descriptor = pd.read_csv("data/Molecular_Descriptor_training.csv")
    era_activity = pd.read_csv("data/ER_activity_training.csv")

    # 移除不必要的列（例如化合物ID）
    X = molecular_descriptor.drop(columns=["SMILES"])
    y = era_activity["pIC50"]

    # 筛选特征
    X_final_filtered = advanced_filter_features(X, y)
